{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exporting ONNX Models with MXNet\n",
    "\n",
    "The [Open Neural Network Exchange](https://onnx.ai/) (ONNX) is an open format for representing deep learning models with an extensible computation graph model, definitions of built-in operators, and standard data types. Starting with MXNet 1.3, models trained using MXNet can now be saved as ONNX models.\n",
    "\n",
    "In this example, we show how to train a model on Amazon SageMaker and save it as an ONNX model. This notebook is based on the [MXNet MNIST notebook](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/mxnet_mnist/mxnet_mnist.ipynb) and the [MXNet example for exporting to ONNX](https://mxnet.incubator.apache.org/tutorials/onnx/export_mxnet_to_onnx.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup\n",
    "\n",
    "First we need to define a few variables that we'll need later in the example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.session import Session\n",
    "\n",
    "# AWS region\n",
    "region = boto3.Session().region_name\n",
    "\n",
    "# S3 bucket for saving code and model artifacts.\n",
    "# Feel free to specify a different bucket here if you wish.\n",
    "bucket = Session().default_bucket()\n",
    "\n",
    "# Location to save your custom code in tar.gz format.\n",
    "custom_code_upload_location = 's3://{}/customcode/mxnet'.format(bucket)\n",
    "\n",
    "# Location where results of model training are saved.\n",
    "model_artifacts_location = 's3://{}/artifacts'.format(bucket)\n",
    "\n",
    "# IAM execution role that gives SageMaker access to resources in your AWS account.\n",
    "# We can use the SageMaker Python SDK to get the role from our notebook environment. \n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The training script\n",
    "\n",
    "The ``mnist.py`` script provides all the code we need for training and hosting a SageMaker model. The script we will use is adaptated from Apache MXNet [MNIST tutorial](https://mxnet.incubator.apache.org/tutorials/python/mnist.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize mnist.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exporting to ONNX\n",
    "\n",
    "The important part of this script can be found in the `save` method. This is where the ONNX model is exported:\n",
    "\n",
    "```python\n",
    "import os\n",
    "\n",
    "from mxnet.contrib import onnx as onnx_mxnet\n",
    "import numpy as np\n",
    "\n",
    "def save(model_dir, model):\n",
    "    symbol_file = os.path.join(model_dir, 'model-symbol.json')\n",
    "    params_file = os.path.join(model_dir, 'model-0000.params')\n",
    "\n",
    "    model.symbol.save(symbol_file)\n",
    "    model.save_params(params_file)\n",
    "\n",
    "    data_shapes = [[dim for dim in data_desc.shape] for data_desc in model.data_shapes]\n",
    "    output_path = os.path.join(model_dir, 'model.onnx')\n",
    "    \n",
    "    onnx_mxnet.export_model(symbol_file, params_file, data_shapes, np.float32, output_path)\n",
    "```\n",
    "\n",
    "The last line in that method, `onnx_mxnet.export_model`, saves the model in the ONNX format. We pass the following arguments:\n",
    "\n",
    "* `symbol_file`: path to the saved input symbol file\n",
    "* `params_file`: path to the saved input params file\n",
    "* `data_shapes`: list of the input shapes\n",
    "* `np.float32`: input data type\n",
    "* `output_path`: path to save the generated ONNX file\n",
    "\n",
    "For more information, see the [MXNet Documentation](https://mxnet.incubator.apache.org/api/python/contrib/onnx.html#mxnet.contrib.onnx.mx2onnx.export_model.export_model)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training the model\n",
    "\n",
    "With the training script written to export an ONNX model, the rest of training process looks like any other Amazon SageMaker training job using MXNet. For a more in-depth explanation of these steps, see the [MXNet MNIST notebook](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/mxnet_mnist/mxnet_mnist.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.mxnet import MXNet\n",
    "\n",
    "mnist_estimator = MXNet(entry_point='mnist.py',\n",
    "                        role=role,\n",
    "                        output_path=model_artifacts_location,\n",
    "                        code_location=custom_code_upload_location,\n",
    "                        train_instance_count=1,\n",
    "                        train_instance_type='ml.m4.xlarge',\n",
    "                        framework_version='1.4.0',\n",
    "                        hyperparameters={'learning-rate': 0.1})\n",
    "\n",
    "train_data_location = 's3://sagemaker-sample-data-{}/mxnet/mnist/train'.format(region)\n",
    "test_data_location = 's3://sagemaker-sample-data-{}/mxnet/mnist/test'.format(region)\n",
    "\n",
    "mnist_estimator.fit({'train': train_data_location, 'test': test_data_location})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next steps\n",
    "\n",
    "Now that we have an ONNX model, we can deploy it to an endpoint in the same way we do in the [MXNet MNIST notebook](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/mxnet_mnist/mxnet_mnist.ipynb).\n",
    "\n",
    "For examples on how to write a `model_fn` to load the ONNX model, please refer to:\n",
    "* the [MXNet ONNX Super Resolution notebook](https://github.com/awslabs/amazon-sagemaker-examples/tree/master/sagemaker-python-sdk/mxnet_onnx_superresolution)\n",
    "* the [MXNet documentation](https://mxnet.incubator.apache.org/api/python/contrib/onnx.html#mxnet.contrib.onnx.onnx2mx.import_model.import_model)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_mxnet_p36",
   "language": "python",
   "name": "conda_mxnet_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
